import numpy as np
import random
import os
from cec17_functions import cec17_test_func
import sys
k = 2#size of the triad k+1
dim = eval(sys.argv[1])#dimension of the solution
NP = eval(sys.argv[2])#the szie of population
stop =35#stagnation times
pm =0.01#restart probability
AS = int(NP/2)#the szie of archive
mean =1.49618#mean of Gaussian sampling
std = 0.1#variance f Gaussian sampling
x_max = 100#maximum boundary of feasible solutions
x_min = -100#minimum boundary of feasible solutions
path = 'data/STTPSO'#The file path where the data is stored
if os.path.exists(path):
    pass
else:
    os.mkdir(path)

def archive_add(location, fitness_X):
    '''his function defines how to add particles to the archive'''
    global num_of_archive,archive
    if num_of_archive < AS:
        #When the particles in the archive are less than the maximum capacity, add location and fitness_X directly
        archive[num_of_archive] = (location, fitness_X)
        num_of_archive += 1
    else:
        #Pick a particle at random and replace it when the particles in the archive are saturated
        old_pbest_position = random.randint(0, AS - 1)
        archive.pop(old_pbest_position)
        archive.append((location, fitness_X))


def init_evalution():
    global new_evaluation,archive
    archive = [0 for _ in range(AS)]#archive of size AS
    location = np.random.uniform(x_min, x_max, (NP, dim))
    #The initial solution is a uniform distribution in the solution space
    Speed = np.zeros((NP, dim))
    #Initial speed is zero
    population_value = np.zeros((NP, 1))
    #Save the fitness value of the current population
    pbest = location.copy()
    #At initialization, pbest is a copy of the population
    gbest = np.zeros(dim)
    #gbest record the current global optimum
    pbestfit = np.zeros(NP)
    #Save the fitness value of pbest
    gbestfit = float('inf')
    #Save the fitness value of gbest

    for i in range(np.size(location, 0)):
        cec17_test_func(location[i], population_value[i], dim, 1, func)
        #Call the cec2017 test set to calculate the fitness value
        archive_add(location[i],pbestfit[i])
        #add to archive
        new_evaluation += 1#Record the number of fitness evaluations
        pbestfit[i] = population_value[i][0] - func * 100
        #Record population fitness and subtract bias terms
        if pbestfit[i] < gbestfit:
            #record the global optimal solution
            gbestfit = pbestfit[i]
            gbest = location[i].copy()
    return gbestfit, archive, pbest, gbest, pbestfit, location, Speed, population_value


def iterator(gbestfit):
    global new_evaluation, gbest,location,population_value,pbestfit,Speed
    meaningless_counter = 0
    #Nothing to do with the algorithm, just to record the current highest fitness value every 5000 evaluations
    tridindex = np.random.randint(0,NP+AS-1, (NP, 2))
    #Create three topology structures for each particle
    stopnumber = [0 for j in range(NP)]
    #Create each particle stagnation counter
    while new_evaluation <= evaluation:
        w = 0.9-new_evaluation/evaluation*0.5
        #Linear decay of inertia weights from 0.9 to 0.4
        for i in range(NP):
            #Gaussian sampling and ensure that c1, c2 are greater than 0
            c1 = np.random.normal(loc=1, scale=std, size=1)[0]
            while c1 <0:
                c1 = np.random.normal(loc=mean, scale=std, size=1)[0]
            c2 = np.random.normal(loc=mean, scale=std, size=1)[0]
            while c2 <0:
                c2 = np.random.normal(loc=mean, scale=std, size=1)[0]
            if c1 >c2:
                #Make sure that c1 is less than c2
                c2,c1 =c1,c2
            best_fit = float('inf')
            #Highest fitness value in triad topology
            sum_X = np.zeros(dim)
            #Mean of all particles in triad topology
            choicelist = [(pbest[i],pbestfit[i])]
            if tridindex[i][0]>=NP:
                choicelist.append(archive[tridindex[i][0]-NP])
            else:
                while tridindex[i][0]==tridindex[i][1] or tridindex[i][0]==i:
                    tridindex[i][0] = random.randint(0 ,NP-1)
                choicelist.append((pbest[tridindex[i][0]],pbestfit[tridindex[i][0]]))
            if tridindex[i][1]>=NP:
                choicelist.append(archive[tridindex[i][1]-NP])
            else:
                while  tridindex[i][1]==i:
                    tridindex[i][1] = random.randint(0 ,NP-1)
                choicelist.append((pbest[tridindex[i][1]],pbestfit[tridindex[i][1]]))
            #Build triad topology with established random neighborhood particle indices
            for choice in choicelist:
                sum_X += choice[0]
                if choice[1] < best_fit:
                    best_fit = choice[1]
                    best_location = choice[0]
            avg_X = sum_X/3
            #Compute the optimal solution and the mean in the topology
            Speed[i] = w* Speed[i] +  c1*np.random.random(dim) * (
                    best_location - location[i]) + c2 * np.random.random(dim) * (
                               avg_X - location[i])
            location[i] = location[i] + Speed[i]
            #Update speed and position by updating formulas
            check = np.where(location[i] > x_max)
            for j in range(len(check[0])):
                location[i][check[0][j]] =x_max
            check = np.where(location[i] < x_min)
            for j in range(len(check[0])):
                location[i][check[0][j]] = x_min
            #Out-of-bounds checking and out-of-bounds processing
            cec17_test_func(location[i], population_value[i], dim, 1, func)
            #Calculate the fitness value of the updated particle
            new_evaluation += 1
            meaningless_counter += 1
            #Record the number of evaluations
            if population_value[i][0] - func * 100 < pbestfit[i]:
                '''When the individual historical optimum changes, 
                add the old pbest to the archive and reset the counter'''
                archive_add(pbest[i], pbestfit[i])
                stopnumber[i]=0
                pbestfit[i] = population_value[i][0] - func * 100
                pbest[i] = location[i].copy()
                if population_value[i][0] - func * 100 < gbestfit:
                    gbest = location[i].copy()
                    gbestfit = population_value[i][0] - func * 100
            else:
                stopnumber[i]+=1
                #After the update, the individual history optimal has not changed, and the stagnation counter is incremented by one.
            if stopnumber[i]>=stop:
                #Reselect neighbor particles when the counter reaches the maximum number of stalls
                tridindex[i][0] = random.randint(0 ,AS+NP-1)
                tridindex[i][1] = random.randint(0, AS+NP-1)

        if random.random() <pm:
            #If the generated random number is less than pm, a new random solution is generated and added to the archive
            data = [0]
            randomlocation = np.random.uniform(x_min, x_max, (dim,))
            cec17_test_func(randomlocation, data, dim, 1, func)
            archive_add(randomlocation,data[0])

        if meaningless_counter >= 5000:
            #The fitness value is recorded every 5000 generations
            meaningless_counter-=5000
            print(gbestfit)
            file.writelines("{:.3e}\t".format(gbestfit))
    file.writelines("{:.3e}\t".format(gbestfit))

evaluation = 10000 * dim
#The maximum number of evaluations is 10000D
for func in range(1,31):
    #The test function serial number is 1-30
    if func != 2:
        #f2 has been deleted by cec2017
        for run in range(30):
            #Repeat 30 independent experiments
            file = open(path + '/dim{}func{}.txt'.format(dim,func), 'a+')
            #Record the results of 30 independent experiments for a function in every file
            num_of_archive = 0
            new_evaluation = 0
            gbestfit, archive, pbest, gbest, pbestfit, location, Speed, population_value= init_evalution()
            iterator(gbestfit)
            file.writelines("\n")
            file.close()
